import matplotlib.pyplot as plt

from pic2video import *


def add_axis0(axis_frames_dir, frames_dir, dpi=50):
    """
    为图片加上坐标轴
    :param axis_frames_dir: 完成坐标轴添加的图片目录
    :param frames_dir: 需要添加的图片目录
    :param dpi: 像素密度，越大越清晰，图片越大
    :return: None
    """
    if not os.path.exists(axis_frames_dir):
        os.makedirs(axis_frames_dir)
    order = [int(i.strip(".jpg")) for i in os.listdir(frames_dir) if i.endswith(".jpg")]
    jpglist = [f"{frames_dir}/{i}.jpg" for i in sorted(order)]  # 直接读取可能非顺序帧

    for i, jpg in enumerate(jpglist):
        image = plt.imread(jpg)
        plt.imshow(image, cmap='gray')
        newjpg = f"{axis_frames_dir}/{os.path.basename(jpg)}"
        plt.savefig(newjpg, dpi=dpi, bbox_inches='tight')
        print(f"进度{(i + 1)}/{len(jpglist)}， 转换带Axis的 {newjpg} 完成！")

    print("对图片添加坐标轴完成！")


if __name__ == "__main__":
    add_axis0("fumo/fumo_axis_frames", frames_dir="fumo/fumo_char_frames",  dpi=50)